import torch
from ultralytics import YOLO
from ultralytics.models.yolo.detect.c2f_transfer import replace_c2f_with_c2f_v2, replace_c2f_v2_with_c2f, replace_c2f_dynamic_with_c2f_dynamic_v2


if __name__ == '__main__':
    input = torch.randn((1, 3, 640, 640))
    
    # C2f -> C2f_V2
    # model_weight_path = r'runs\train\yolov8base\weights\best.pt'
    # model = YOLO(model_weight_path)
    # model.model.eval()
    # pre_res = model.model(input)[0]
    # replace_c2f_with_c2f_v2(model.model)
    # model.model.eval()
    # after_res = model.model(input)[0]
    # print(torch.mean(pre_res - after_res))
    # torch.save({'model':model.model.half()}, f'{model_weight_path[:model_weight_path.rfind(".")]}_v2.pt')
    
    # # C2f_V2 -> C2f
    # model_v2_weight_path = 'runs/train/exp/weights/best.pt'
    # model = YOLO(model_v2_weight_path)
    # model.model.eval()
    # pre_res = model.model(input)[0]
    # replace_c2f_v2_with_c2f(model.model)
    # model.model.eval()
    # after_res = model.model(input)[0]
    # print(torch.mean(pre_res - after_res))
    # torch.save({'model':model.model.half()}, f'{model_v2_weight_path[:model_v2_weight_path.rfind(".")]}_notv2.pt')


    # C2f_DynamicAxialRouterV2 -> C2f_DynamicAxialRouterV2_v2
    model_weight_path = r'runs\train\yolov8-c2famhabi\weights\best.pt'
    model = YOLO(model_weight_path)  # 加载模型权重
    model.model.eval()
    # 预转换时前向传播输出结果
    pre_res = model.model(input)[0]
    # 替换所有 C2f_DynamicAxialRouterV2 为 C2f_DynamicAxialRouterV2_v2 并转换权重
    replace_c2f_dynamic_with_c2f_dynamic_v2(model.model)
    model.model.eval()
    # 转换后前向传播输出结果
    after_res = model.model(input)[0]
    # 比较前后输出的均值差异
    print(torch.mean(pre_res - after_res))
    # 保存转换后的模型（注意这里使用 .half() 根据你的需求转换精度）
    torch.save({'model': model.model}, f'{model_weight_path[:model_weight_path.rfind(".")]}_v2.pt')